import re

from hhutil.io import read_lines, eglob
from hhutil.functools import lmap

import numpy as np


epoch_p = re.compile(r"""Epoch \d+/\d+, lr: (\d\.\d{6})""")
def get_lr(s):
    return float(epoch_p.search(s).group(1))

train_p = re.compile(r""".* Train \d+/\d+ - loss: (\d\.\d{4}) - acc: (\d\.\d{4})""")
def extract_train(s):
    m = train_p.search(s)
    return float(m.group(1)), float(m.group(2))

valid_p = re.compile(r""".* Valid \d+/\d+ - loss: (\d\.\d{4}) - acc: (\d\.\d{4})""")
def extract_valid(s):
    m = valid_p.search(s)
    return float(m.group(1)), float(m.group(2))

def extract(fp):
    lines = read_lines(fp)
    epoch_lines = []
    train_lines = []
    valid_lines = []

    i = 0
    while not lines[i].startswith("Epoch 1/"):
        i += 1

    lines = lines[i + 1:]
    for l in lines:
        if 'Epoch' in l:
            epoch_lines.append(l)
        elif 'Train' in l:
            train_lines.append(l)
        elif 'Valid' in l:
            valid_lines.append(l)


    lrs = lmap(get_lr, epoch_lines[1:])

    train_losses, train_accs = zip(*map(extract_train, train_lines))
    valid_losses, valid_accs = zip(*map(extract_valid, valid_lines))

    lrs, train_losses, train_accs, valid_losses, valid_accs = [
        np.array(x) for x in [lrs, train_losses, train_accs, valid_losses, valid_accs]]
    return lrs, train_losses, train_accs, valid_losses, valid_accs

def get_valid_accs(f):
    return extract(f)[4]

log_dir = "/Users/hrvvi/Downloads/PyramidNet2"

log_files1 = list(eglob(log_dir, "4-*.log"))
accs1 = []
for f in log_files1:
    accs1.append(np.median(get_valid_accs(f)[-8:]))

log_files2 = list(eglob(log_dir, "13-*.log"))
accs2 = []
for f in log_files2:
    accs2.append(np.median(get_valid_accs(f)[-8:]))

print(np.median(accs1), np.median(accs2))
# errs = (1 - valid_accs) * 100
# plt.plot(errs)
# plt.ylim([-0.9, 20])
#
# errs = (1 - train_accs) * 100
# plt.plot(errs)
# plt.ylim([-0.9, 20])